;;;; -*- Mode: Lisp; -*-

(in-package :lib.gl)

;;; ------------------------------------------------- 
;;; Matrix math using quaternions.  Liberally borrowed
;;; from GamePhysics.
;;;

(defun quaternion-p (v)
  (and (or (typep v 'external-vector) (typep v 'vector))
       (= (vector-length v) 4)))

(deftype quaternion () `(satisfies quaternion-p))

(defun quat-w (q)
  (mref q 0))

(defun quat-x (q)
  (mref q 1))

(defun quat-y (q)
  (mref q 2))

(defun quat-z (q)
  (mref q 3))

(defun make-quat (w x y z)
  (vector w x y z))

(defun copy-quat (q)
  (make-quat (quat-w q)
	     (quat-x q)
	     (quat-y q)
	     (quat-z q)))

(defun quat-conjugate (q)
  (make-quat (quat-w q) (- (quat-x q)) (- (quat-y q)) (- (quat-z q))))

(defun quat-magnitude (q)
  (let ((ftx (quat-x q))
        (fty (quat-y q))
        (ftz (quat-z q))
        (ftw (quat-w q)))
    (sqrt (+ (* ftx ftx) (* fty fty) (* ftz ftz) (* ftw ftw)))))

(defun quat-mult (q1 q2)
  (nquat-mult q1 (copy-quat q2)))

(defun %nquat-mult (q1 q2)
  (let ((x1 (quat-x q1)) (y1 (quat-y q1)) (z1 (quat-z q1)) (w1 (quat-w q1))
        (x2 (quat-x q2)) (y2 (quat-y q2)) (z2 (quat-z q2)) (w2 (quat-w q2)))  
    (setf (mref q2 0) (- (* w2 w1) (* x2 x1)    (* y2 y1)    (* z2 z1))
	  (mref q2 1) (+ (* w2 x1) (* x2 w1)    (* -1 y2 z1) (* z2 y1))
	  (mref q2 2) (+ (* w2 y1) (* x2 z1)    (* y2 w1)    (* -1 z2 x1))
	  (mref q2 3) (+ (* w2 z1) (* -1 x2 y1) (* y2 x1)    (* z2 w1)))

    ;; (setf (mref q2 0) (- (* w1 w2) (* x1 x2)    (* y1 y2)    (* z1 z2))
    ;; 	  (mref q2 1) (+ (* w1 x2) (* x1 w2)    (* y1 z2)    (* -1 z1 y2))
    ;; 	  (mref q2 2) (+ (* w1 y2) (* -1 x1 z2) (* y1 w2)    (* z1 x2))
    ;; 	  (mref q2 3) (+ (* w1 z2) (* x1 y2)    (* -1 y1 x2) (* z1 w2)))
    q2))

(defparameter *quat-reset-limit* 10.0)

(defun nquat-mult (q1 q2)
  (restart-case
      (handler-case (%nquat-mult q1 q2)
	(floating-point-overflow () (invoke-restart 'normalize-first)))
    (normalize-first () (progn 
			  (normalize q1)
			  (normalize q2)
			  (%nquat-mult q1 q2))))
  (when (> (abs (quat-w q2)) *quat-reset-limit*)
    ;(format t "Normalizing quaternion ~A." q2)
    (normalize q2))
  q2)

(defun quat-vector-mult (q1 vec)
  (quat-mult q1 (vector-to-quat vec)))
    
(defun %euler-to-quat-simple (angle-xx angle-yy angle-zz)
  (let ((qx nil)
	(qy nil)
	(qz nil))
    (unless (zerop angle-xx)
      (setq qx (axis-angle-to-quat (unit-x-vector) angle-xx)))
    (unless (zerop angle-yy)
      (setq qy (axis-angle-to-quat (unit-y-vector) angle-yy)))
    (unless (zerop angle-zz)
      (setq qz (axis-angle-to-quat (unit-z-vector) angle-zz)))
    (cond ((and (null qx) (null qy) (null qz)) (axis-angle-to-quat (unit-x-vector) 0.0))
	  ((and qx (null qy) (null qz)) qx)
	  ((and qy (null qx) (null qz)) qy)
	  ((and qz (null qx) (null qy)) qz)
	  ((and qx qy (null qz)) (nquat-mult qx qy))
	  ((and qx qz (null qy)) (nquat-mult qx qz))
	  ((and qy qz (null qx)) (nquat-mult qy qz))
	  (t (nquat-mult qx (nquat-mult qy qz))))))

(defun euler-to-quat (angle-xx angle-yy angle-zz)
  (%euler-to-quat-fast angle-xx angle-yy angle-zz))

(defun euler-rotate-matrix (rot-x rot-y rot-z &optional (in-radians t))
 ;;; For rotation matrices
  (when (not in-radians)
    (setf rot-x (* 0.5 rot-x +PI/180+)
	  rot-y (* 0.5 rot-y +PI/180+)
	  rot-z (* 0.5 rot-z +PI/180+)))
  (let ((r-xy (matrix3-by-rows (make-vector (cos rot-z) (- (sin rot-z)) 0.0)
			       (make-vector (sin rot-z) (cos rot-z)     0.0)
			       (make-vector 0.0         0.0             1.0)))
	(r-xz (matrix3-by-rows (make-vector (cos rot-y)     0.0 (sin rot-y))
			       (make-vector 0.0             1.0 0.0)
			       (make-vector (- (sin rot-y)) 0.0 (cos rot-y))))
	(r-yz (matrix3-by-rows (make-vector 1.0 0.0         0.0)
			       (make-vector 0.0 (cos rot-x) (- (sin rot-x)))
			       (make-vector 0.0 (sin rot-x) (cos rot-x)))))
    (m* r-yz r-xz r-xy)))


(defun %euler-to-quat-fast (angle-xx angle-yy angle-zz)
  "Calculate the quaternion for the rotation given by euler angles (in radians) about the 
x, y, and z axes."
  (let* ((r (/ angle-xx 2.0)) (p (/ angle-yy 2.0)) (y (/ angle-zz 2.0)) ; Pitch, yaw, roll
	 (sinp (sin p)) (cosp (cos p))
	 (siny (sin y)) (cosy (cos y))
	 (sinr (sin r)) (cosr (cos r)))
    (make-quat (+ (* cosr cosp cosy) (* sinr sinp siny))
	       (- (* sinr cosp cosy) (* cosr sinp siny))
	       (+ (* cosr sinp cosy) (* sinr cosp siny))
	       (- (* cosr cosp siny) (* sinr sinp cosy)))))
		

(defun axis-angle-to-quat (vec angle)
  "Angle is in radians."
  (let ((c (cos (/ angle 2)))
        (s (sin (/ angle 2))))
    (make-quat c
               (* (x vec) s)
               (* (y vec) s)
               (* (z vec) s))))

(defun quat-to-axis-angle (quat)
  (when (> (quat-w quat) +one+)
    (normalize quat))
  (if (< (abs (- (quat-w quat) +one+)) +epsilon+)
      (values +zero+ (unit-x-vector))
      (let* ((angle (acos (quat-w quat)))
	     (n (/ 1 (sqrt (- 1 (expt (quat-w quat) 2))))))
	(values (* angle +two+) (make-vector (* n (quat-x quat))
					     (* n (quat-y quat))
					   (* n (quat-z quat)))))))
				  
(defun vector-to-quat (vec)
  "Convert 3-dimensional vector vec to a quaternion.  Assumes vec is an euler rotation of 
0 degress about the axis defined by vec."
  (axis-angle-to-quat vec +zero+))

(defun %quat-to-rotation-matrix (v &optional (rm nil))
  (unless (typep v 'quaternion)
    (error "Vector ~a is not a quaternion and cannot be converted to a rotation matrix." v))
  (let* ((rot (or rm (zero-matrix 3)))
	 (w (quat-w v))
	 (x (quat-x v))
	 (y (quat-y v))
	 (z (quat-z v))
	 (xx (* x x))
	 (yy (* y y))
	 (zz (* z z)))
    (setf (mref rot 0 0) (- 1 (* 2 yy) (* 2 zz))
	  (mref rot 0 1) (- (* 2 x y) (* 2 w z))
	  (mref rot 0 2) (+ (* 2 x z) (* 2 w y))
	  (mref rot 1 0) (+ (* 2 x y) (* 2 w z))
	  (mref rot 1 1) (- 1 (* 2 xx) (* 2 zz))
	  (mref rot 1 2) (- (* 2 y z) (* 2 w x))
	  (mref rot 2 0) (- (* 2 x z) (* 2 w y))
	  (mref rot 2 1) (+ (* 2 y z) (* 2 w x))
	  (mref rot 2 2) (- 1 (* 2 xx) (* 2 yy)))
    rot))

(defmethod quat-to-rotation-matrix ((v vector) &optional (rm nil))
  (when (> (abs (quat-w v)) *quat-reset-limit*)
    ;(format t "Normalizing quaternion ~A.~%" v)
    (normalize v))
  (restart-case
      (handler-case (%quat-to-rotation-matrix v)
	(floating-point-overflow () (invoke-restart 'normalize-first)))
    (normalize-first () (progn 
			  (warn "Overflow on quaternion ~A, trying to normalize!" v)
			  (m/ v (quat-w v))
			  (%quat-to-rotation-matrix v rm)))))

(defmethod quat-to-homogeneous-matrix ((v vector) &optional (rs-matrix nil))
  (unless (typep v 'quaternion)
    (error "Vector ~a is not a quaternion and cannot be converted to a rotation matrix." v))
  (let* ((rot (or rs-matrix (zero-matrix 4)))
	 (w (quat-w v))
	 (x (quat-x v))
	 (y (quat-y v))
	 (z (quat-z v))
	 (xx (* x x))
	 (yy (* y y))
	 (zz (* z z)))
    (setf (mref rot 0 0) (- 1 (* 2 yy) (* 2 zz))
	  (mref rot 0 1) (- (* 2 x y) (* 2 w z))
	  (mref rot 0 2) (+ (* 2 x z) (* 2 w y))
	  (mref rot 1 0) (+ (* 2 x y) (* 2 w z))
	  (mref rot 1 1) (- 1 (* 2 xx) (* 2 zz))
	  (mref rot 1 2) (- (* 2 y z) (* 2 w x))
	  (mref rot 2 0) (- (* 2 x z) (* 2 w y))
	  (mref rot 2 1) (+ (* 2 y z) (* 2 w x))
	  (mref rot 2 2) (- 1 (* 2 xx) (* 2 yy))
	  (mref rot 3 3) 1.0)
    rot))

(defmethod quat-from-rotation-matrix ((rot array))
  "Algorithm to convert a rotation matrix into a quaternion, given in Ken Shoemake's article in 1987 SIGGRAPH course notes \"Quaternion Calculus and Fast Animation\"."
  (let ((ftrace (+ (mref rot 0 0) (mref rot 1 1) (mref rot 2 2)))
        (q (make-vector +zero+ +zero+ +zero+ +zero+))
        (ms-inext (make-vector 1 2 0))
        froot)
    (if (> ftrace 0)
        ;; |w| > 1/2, may as well choose w > 1/2
        (progn
          (setq froot (sqrt (+ ftrace 1.0))) ; 2w
          (setf (mref q 0) (* 0.5 froot))
          (setq froot (/ 0.5 froot)) ; 1/(4w)
          (setf (mref q 1) (* froot (- (mref rot 2 1)
                                        (mref rot 1 2))))
          (setf (mref q 2) (* froot (- (mref rot 0 2)
                                        (mref rot 2 0))))
          (setf (mref q 3) (* froot (- (mref rot 1 0)
                                        (mref rot 0 1)))))
      (let ((i 0))
        ;; |w| <= 1/2
        (if (> (mref rot 1 1) (mref rot 0 0))
            (setq i 1))
        (if (> (mref rot 2 2) (mref rot i i))
            (setq i 2))
        (let* ((j (mref ms-inext i))
               (k (mref ms-inext j)))
          (setq froot (sqrt (+ 1.0 
                               (- (mref rot i i)
                                  (mref rot j j)
                                  (mref rot k k)))))
          (setf (mref q (1+ i)) (* 0.5 froot))
          (setq froot (/ 0.5 froot))
          (setf (mref q 0) (* froot (- (mref rot k j) (mref rot j k)))
                (mref q (1+ j)) (* froot (+ (mref rot j i) (mref rot i j)))
                (mref q (1+ k)) (* froot (+ (mref rot k i) (mref rot i k)))))))
    q))




(defun rotation-matrix-to-quat (rot)
  (quat-from-rotation-matrix rot))


(defun slerp-quaternion (tm p q)
  (when (and (quaternion-p p) (quaternion-p q))
    (let* ((fcos (dot-product p q))
           (angle (acos fcos)))
      (if (> (abs angle) SINGLE-FLOAT-EPSILON)
          p
        (let* ((fsin (sin angle))
               (invsin (/ 1.0 fsin))
               (coeff0 (* invsin (sin (* (- 1.0 tm) angle))))
               (coeff1 (* invsin (sin (* tm angle)))))
          (m+ (m* coeff0 p) (m* coeff1 q)))))))

;;;; tests

(defun test-euler-1 ()
  (let ((q1 (euler-to-quat 3.14159 0.0 0.0))
	(q2 (euler-to-quat 0.0 3.14159 0.0))
	(q3 (euler-to-quat 0.0 0.0 (/ 3.14159 2.0))))
    (multiple-value-bind (a v) (quat-to-axis-angle q1)
      (format t "~A ~A ~A~%" q1 a v)
      (assert (< (- a 3.14159) 0.001))
      (assert (vector= v (make-vector 1.0 0.0 0.0))))
    (multiple-value-bind (a v) (quat-to-axis-angle q2)
      (format t "~A ~A ~A~%" q2 a v)
      (assert (< (- a 3.14159) 0.001))
      (assert (vector= v (make-vector 0.0 1.0 0.0))))
    (multiple-value-bind (a v) (quat-to-axis-angle q3)
      (format t "~A ~A ~A~%" q3 a v)
      (assert (< (- a (/ 3.14159 2.0)) 0.001))
      (assert (vector= v (make-vector 0.0 0.0 1.0))))))


(defun test-euler-2 ()
  (dotimes (i 1000)
    (euler-to-quat 0.5 0.5 0.5)))

(defun test-euler-3 ()
  (let ((q1 (euler-to-quat 0.0 0.0 (/ 3.14159 2.0)))
	(q2 (axis-angle-to-quat (unit-z-vector) (/ 3.14159 2.0)))
	(r1 (euler-rotate-matrix 0.0 0.0 (/ 3.14159 2.0)))
	(p (unit-x-vector)))
    (format t "~A ~A~%" q1 q2)
    (format t "~A~%" (m* r1 p))
    (format t "~A~%" (quat-mult (quat-vector-mult q1 p) (quat-conjugate q1)))
    (format t "~A~%" (m* (quat-to-rotation-matrix q1) p))
    (format t "~A~%" (m* (quat-to-rotation-matrix q2) p))
    (format t "~A~%" (magnitude (m- (m* r1 p) (m* (quat-to-rotation-matrix q1) p))))
    (assert (vector= q1 q2))
    (assert (< (magnitude (m- (m* r1 p) (m* (quat-to-rotation-matrix q1) p))) +epsilon+))
    (assert (< (magnitude (m- (m* r1 p) (m* (quat-to-rotation-matrix q2) p))) +epsilon+))))

(defun test-euler-4 ()
  (let ((q1 (euler-to-quat 0.0 (/ 3.14159 2.0) 0.0))
	(q2 (axis-angle-to-quat (unit-y-vector) (/ 3.14159 2.0)))
	(r1 (euler-rotate-matrix 0.0 (/ 3.14159 2.0) 0.0))
	(p (unit-x-vector)))
    (format t "~A ~A~%" q1 q2)
    (format t "~A~%" (m* r1 p))
    (format t "~A~%" (quat-mult (quat-vector-mult q1 p) (quat-conjugate q1)))
    (format t "~A~%" (m* (quat-to-rotation-matrix q1) p))
    (format t "~A~%" (m* (quat-to-rotation-matrix q2) p))
    (format t "~A~%" (magnitude (m- (m* r1 p) (m* (quat-to-rotation-matrix q1) p))))
    (assert (vector= q1 q2))
    (assert (< (magnitude (m- (m* r1 p) (m* (quat-to-rotation-matrix q1) p))) +epsilon+))
    (assert (< (magnitude (m- (m* r1 p) (m* (quat-to-rotation-matrix q2) p))) +epsilon+))))

(defun test-euler-overflow ()
  (let ((q1 (euler-to-quat 0.0 0.1 0.0))
	(p (unit-x-vector)))
    (dotimes (i 126)
      (format t "~A:~A:~A~%" q1 (m* (quat-to-rotation-matrix q1) p) (quat-magnitude q1))
      (setf q1 (quat-mult q1 q1)))
    (print-matrix (quat-to-rotation-matrix q1))))

(defun test-euler-performance ()
  (time (dotimes (i 1000)
	  (%euler-to-quat-simple 0.5 0.0 0.0)))
  (time (dotimes (i 1000)
	  (%euler-to-quat-simple 0.5 0.5 0.0)))
  (time (dotimes (i 1000)
	  (%euler-to-quat-simple 0.5 0.5 0.5)))
  (time (dotimes (i 1000)
	  (%euler-to-quat-fast 0.5 0.5 0.0))))
      
(defun test-euler-performance-2 ()
  (let ((m (zero-matrix 4))
	(q (euler-to-quat 1.7 0.0 0.0)))
    (time (dotimes (i 1000)
	    (quat-to-homogeneous-matrix q)))
    (time (dotimes (i 1000)
	    (quat-to-homogeneous-matrix q m)))))
